# -*- coding: utf-8 -*-

"""
    2025-01-30 by Jie Zheng
"""


import numpy as np
import astropy.io.fits as fits
from astropy.wcs.utils import fit_wcs_from_points
from astropy.coordinates import SkyCoord as Coord
import astroalign as aa
import qmatch as qm
from .wcs_util import *


def cetian(sx, sy, smag, rra, rdec, rmag, crpix, crval, scount=20, rcount=60, sip_deg=3):
    """
    WCS processing
    :param sx: x coordinate of source image
    :param sy: y coordinate of source image
    :param smag: magnitude of source image
    :param rra: ra coordinate of reference image
    :param rdec: dec coordinate of reference image
    :param rmag: magnitude of reference image
    :param crpix: crpix of source image, usually the center of the image
    :param crval: crval of source image, usually the guess of the center of the image
    :param scount: count of source stars used
    :param rcount: count of reference stars used
    :param sip_deg: sip degree
    :return: wcs object, and residual
    select brightest stars from image stars and reference stars, find the linear wcs by aa,
    then match image stars to reference stars, find the sip coefficients
    """

    # sort stars by magnitude, and select brightest stars
    ixs = np.argsort(smag)
    ixr = np.argsort(rmag)
    rrab, rdecb, rmagb = rra[ixr][:rcount], rdec[ixr][:rcount], rmag[ixr][:rcount]
    sxb, syb, smagb = sx[ixs][:scount], sy[ixs][:scount], smag[ixs][:scount]

    # alpha,delta --> xsi,eta
    rxsib, retab = ad2xe(rrab, rdecb, crval)
    rxsib *= 3600.0
    retab *= 3600.0
    # x,y --> u,v
    sub, svb = xy2uv(sxb, syb, crpix)

    # match with aa
    tr, (suvb, rxeb) = aa.find_transform(np.c_[sub, svb], np.c_[rxsib, retab])
    
    # find the new crval
    (xsi1, eta1), = tr([[0,0]])
    ra1, dec1 = xe2ad(np.array([xsi1])/3600, np.array([eta1])/3600, crval)
    crval1 = (ra1[0], dec1[0])

    # star u,v to xsi,eta
    sxeb = tr(np.c_[sub, svb])
    sxsib, setab = sxeb.T

    # match star with reference, by xsi,eta
    sixb, rixb = qm.match2d(sxsib, setab, rxsib, retab, )

    # fit with astropy
    ss = fit_wcs_from_points(
        (sxb[sixb], syb[sixb]),
        Coord(rrab[rixb], rdecb[rixb], frame="fk5", unit="deg"),
        proj_point=Coord(*crval1, frame="fk5", unit="deg"),
        sip_degree=sip_deg,
    )

    # star x,y to ra,dec
    scoo = ss.pixel_to_world(sx, sy)
    sra = scoo.ra.deg
    sdec = scoo.dec.deg
    # match star with reference, by ra,dec
    six, rix = qm.match_radec(sra, sdec, rra, rdec, dislimit=0.001)
    # residual
    res_ra  = (sra [six] - rra [rix]) * 3600
    res_dec = (sdec[six] - rdec[rix]) * 3600

    return ss, res_ra, res_dec
